from contextlib import asynccontextmanager
from typing import Any
import mcp.server.stdio
from collections.abc import AsyncIterator
import json
import mcp.types as types
from mcp.server.lowlevel import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import asyncio


# 模拟一个数据库管理类
class MockDatabase:
    # 类方法 用于异步连接数据库，返回数据库的实例
    @classmethod
    async def connect(cls) -> "MockDatabase":
        print("数据库已经连接")
        return cls()

    async def disconnect(self) -> None:
        print("数据库已经断开")

    async def query(self, query_str: str) -> list[dict[str, Any]]:
        return [{"id": 1, "name": "张三", "query": query_str}]


# 2.定义生命周期管理函数，负责资源的初始化和清理
@asynccontextmanager
async def server_lifespan(_server: Server) -> AsyncIterator[dict[str, Any]]:
    # 在启动的时候连接数据库资源
    db = await MockDatabase.connect()
    try:
        # 将数据库的资源传递给请求处理器
        yield {"db": db}
    finally:
        await db.disconnect()


# 3.创建低层Server实例，传入生命周期管理函数
server = Server("lowlevel-server", lifespan=server_lifespan)


# 4.注册工具列表
@server.list_tools()
async def handle_list_tool() -> list[types.Tool]:
    return [
        types.Tool(
            name="query_db",  # 工具的名称
            description="查询数据库",  # 工具的描述
            inputSchema={  # 输入参数
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "要执行的SQL语句"}
                },
                "required": ["query"],
            },
            outputSchema={  # 结构化输出
                "type": "object",
                "properties": {
                    "results": {
                        "type": "array",
                        "items": {
                            "type": "object",
                            "properties": {
                                "id": {"type": "integer"},
                                "name": {"type": "string"},
                                "query": {"type": "string"},
                            },
                        },
                    },
                    "count": {"type": "integer"},
                },
            },
        )
    ]


# 5.注册工具调用处理器
@server.call_tool()
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]:
    if name != "query_db":
        raise ValueError(f"未知工具:{name}")
    # 获取生命周期上下文对象
    ctx = server.request_context
    # lifespan_context里面有就有数据库管理器db
    db = ctx.lifespan_context["db"]
    query_str = arguments["query"]
    results = await db.query(query_str)
    return {"results": results, "count": len(results)}


# 6.注册资源列表处理器
@server.list_resources()
async def handle_list_resources() -> list[types.Resource]:
    return [
        types.Resource(
            uri="data://status",  # 资源的URI
            name="系统状态",  # 资源名称
            description="获取系统的运行状态",
            mimeType="application/json",  # MIME类型
        )
    ]


# 7.注册资源读取处理器，处理资源读取请求
@server.read_resource()
async def handle_read_resource(uri: str) -> list[types.TextContent]:
    if str(uri) == "data://status":
        return json.dumps({"status": "running"}, indent=2)
    else:
        raise ValueError(f"未知资源:{uri}")


async def run():
    # 使用stdio传输方式启动服务器
    async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            InitializationOptions(
                server_name="lowlevel-server",
                server_version="0.1.0",
                capabilities=server.get_capabilities(
                    notification_options=NotificationOptions(),
                    experimental_capabilities={},
                ),
            ),
        )


if __name__ == "__main__":
    asyncio.run(run())
